#!/usr/bin/python

import sys,os,re
from pylab import *
from scipy.ndimage import filters,morphology,interpolation,measurements
from scipy import stats
import tables
from llpy.sutils import *

import argparse
parser = argparse.ArgumentParser(description="Shuffle the samples in a file.")
parser.add_argument('input',default='data.h5',help="data file")
parser.add_argument('-d','--size',default=24,type=int,help='target size')
parser.add_argument('-o','--output',default='zoomed.h5',help="output file")
parser.add_argument('-D','--display',action="store_true",help="debugging display")
parser.add_argument('-N','--nsamples',type=int,default=2000000000,help="max # of samples")
parser.add_argument('-r','--pattern',default='.*',help='pattern selection (regex)')
args = parser.parse_args()

scale = None

select = re.compile("^"+args.pattern+"$")

with tables.openFile(args.input) as db, tables.openFile(args.output,"w") as odb:
    table_lcopy(db,odb)
    table_log(odb,str(sys.argv))

    nsamples = len(db.root.classes)
    classes = db.root.classes[:nsamples]
    shape = list(db.root.patches[0].shape)
    strs = [udecode(c) for c in classes]
    print "original samples",nsamples
    samples = find([select.match(s) is not None for s in strs])
    print "matching samples",len(samples)
    random.shuffle(samples)
    samples = samples[:min(nsamples,args.nsamples)]

    rel = db.root.rel if "rel" in dir(db.root) else None
    odb.createEArray(odb.root,"patches",tables.Float32Atom(),shape=[0]+shape,filters=tables.Filters(9))
    odb.createEArray(odb.root,"classes",tables.Int64Atom(),shape=(0,),filters=tables.Filters(9))
    if rel is not None:
        odb.createEArray(odb.root,"rel",tables.Float32Atom(),shape=(0,3),filters=tables.Filters(9))
    for i in samples:
        n = len(odb.root.classes)
        if n%1000==0: print n
        p = db.root.patches[i]
        if args.display:
            ion(); clf(); gray(); imshow(p,interpolation='nearest'); ginput(1,0.001)
        odb.root.classes.append([db.root.classes[i]])
        odb.root.patches.append([p])
        if rel is not None:
            odb.root.rel.append([rel[i]])
